package similarity

import "math"

// documentSet 文档结构体
type documentSet struct {
	Document [][]string
	IDF      map[string]float64
	TF       []map[string]float64
}

// calculateTF 计算词频
func (d *documentSet) calculateTF() {
	for _, words := range d.Document {
		tf := make(map[string]float64)
		totalWords := float64(len(words))

		// 统计词频
		wordCount := make(map[string]int)
		for _, word := range words {
			wordCount[word]++
		}

		// 计算TF值
		for word, count := range wordCount {
			tf[word] = float64(count) / totalWords
		}
		d.TF = append(d.TF, tf)
	}
}

// calculateIDF 计算逆文档频率
func (d *documentSet) calculateIDF() {
	idf := make(map[string]float64)
	totalDocs := float64(len(d.Document))

	// 统计包含每个词的文档数
	wordInDocs := make(map[string]int)
	for _, doc := range d.Document {
		// 使用map去重，确保每篇文档中同一个词只计算一次
		uniqueWords := make(map[string]struct{})
		for _, word := range doc {
			uniqueWords[word] = struct{}{}
		}

		for word := range uniqueWords {
			wordInDocs[word]++
		}
	}

	// 计算IDF值
	for word, docCount := range wordInDocs {
		idf[word] = math.Log(totalDocs / float64(docCount))
	}

	d.IDF = idf
}

// calculateTFIDF 计算TF-IDF值
func (d *documentSet) calculateTFIDF() []map[string]float64 {
	var result []map[string]float64
	for _, tf := range d.TF {
		tfidf := make(map[string]float64)
		for word, tfValue := range tf {
			idfValue, exists := d.IDF[word]
			if !exists {
				idfValue = 0
			}
			tfidf[word] = tfValue * idfValue
		}
		result = append(result, tfidf)
	}

	return result
}
